package com.mixed.mixedserver.config;

import com.mixed.mixedserver.config.listeners.TestChatListener;
import com.mixed.mixedserver.service.ai.*;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.model.openai.OpenAiTokenCountEstimator;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import jakarta.annotation.Resource;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class LLMConfig {

    @Resource
    private RedisChatMemoryStore redisChatMemoryStore;


    @Bean("QWen")
    public ChatModel  QWen() {
        OpenAiChatModel build = OpenAiChatModel.builder()
                .apiKey(System.getenv("QWen-key"))
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .logRequests( true)
                .logResponses( true)
                .listeners(List.of(new TestChatListener()))
                .modelName("qwen-plus").build();
        return build;
    }

    @Bean("DeepSeek")
    public ChatModel  DeepSeek() {
        OpenAiChatModel build = OpenAiChatModel.builder()
                .apiKey(System.getenv("DeepSeek-key"))
                .baseUrl("https://api.deepseek.com/v1")
                .modelName("deepseek-reasoner")
                .logRequests( true)
                .logResponses( true)
                .listeners(List.of(new TestChatListener()))
                // 移除responseFormat以避免与流式传输冲突
//                .responseFormat("json_object")
//                .maxTokens(8192)
                .build();
        return build;
    }

    @Bean("DeepSeekStreaming")
    public StreamingChatModel DeepSeekStreaming() {
        OpenAiStreamingChatModel build = OpenAiStreamingChatModel.builder()
                .apiKey(System.getenv("DeepSeek-key"))
                .baseUrl("https://api.deepseek.com/v1")
                .modelName("deepseek-reasoner")
                .logRequests( true)
                .logResponses( true)
                .listeners(List.of(new TestChatListener()))
//                .responseFormat("json_object")
//                .maxTokens(8192)
                .build();
                return build;
    }

    @Bean
    public ChatAssistant chatAssistant(@Qualifier("QWen") ChatModel chatModel){
        ChatAssistant assistant = AiServices.create(ChatAssistant.class, chatModel);
        return assistant;
    }

    @Bean("QWenImageModel")
    public ChatModel qWenImageModel(){
        ChatModel build = OpenAiChatModel.builder()
                .modelName("qwen-vl-max")
                .apiKey(System.getenv("QWen-key"))
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .logRequests( true)
                .logResponses( true)
                .listeners(List.of(new TestChatListener()))
                .build();
        return build;
    }

    @Bean
    public StreamingChatModel streamingChatModel(){
        StreamingChatModel build = OpenAiStreamingChatModel.builder()
                .apiKey(System.getenv("QWen-key"))
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .modelName("qwen-plus")
                .build();
        return build;
    }

    @Bean("chatAssistantFlux")
    public ChatAssistant chatAssistantFlux(@Qualifier("streamingChatModel") StreamingChatModel chatModel){
        ChatAssistant assistant = AiServices.create(ChatAssistant.class, chatModel);
        return assistant;
    }

    @Bean(name = "chatMessageWindowChatMemory")
    public ChatMemoryAssistant chatMessageWindowChatMemory(@Qualifier("QWen") ChatModel chatModel){
        ChatMemoryAssistant assistant = AiServices
                .builder(ChatMemoryAssistant.class)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(100))
                .chatModel(chatModel).build();
        return assistant;
    }


    @Bean(name = "chatTokenWindowChatMemory")
    public ChatMemoryAssistant chatTokenWindowChatMemory(@Qualifier("QWen") ChatModel chatModel){
        ChatMemoryAssistant assistant = AiServices
                .builder(ChatMemoryAssistant.class)
                .chatMemoryProvider(memoryId -> TokenWindowChatMemory.withMaxTokens(2000,new OpenAiTokenCountEstimator("gpt-4")))
                .chatModel(chatModel).build();
        return assistant;
    }

    @Bean(name = "lawChatAssistant")
    public LawChatAssistant lawChatAssistant(@Qualifier("QWen") ChatModel chatModel){
        return AiServices.create(LawChatAssistant.class, chatModel);
    }

    @Bean
    public ChatPersistenceAssistant chatPersistenceAssistant(@Qualifier("QWen") ChatModel chatModel){
        ChatMemoryProvider chatMemoryProvider  = memoryId -> MessageWindowChatMemory
                .builder()
                .id(memoryId)
                .maxMessages(1000)
                .chatMemoryStore(redisChatMemoryStore)
                .build();
        return AiServices.builder(ChatPersistenceAssistant.class)
                .chatModel(chatModel)
                .chatMemoryProvider(chatMemoryProvider)
                .build();
    }

    @Bean()
    public OpenAiEmbeddingModel openAiEmbeddingModel() {
        OpenAiEmbeddingModel build = OpenAiEmbeddingModel.builder()
                .apiKey(System.getenv("QWen-key"))
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                .modelName("text-embedding-v3")
                .build();
        return build;
    }

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore(OpenAiEmbeddingModel embeddingModel ) {
        EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.builder()
                .host("abluedog.fun")                           // Required: Host of the PostgreSQL instance
                .port(5432)                                  // Required: Port of the PostgreSQL instance
                .database("vector_db")                        // Required: Database name
                .user("abluedog")                             // Required: Database user
                .password("lys2495150633")                     // Required: Database password
                .table("documents")                      // Required: Table name to store embeddings
                .dimension(embeddingModel.dimension())       // Required: Dimension of embeddings
                .build();
        return embeddingStore;
    }

    @Bean
    public InMemoryEmbeddingStore<TextSegment> inMemoryChatMemoryStore() {
        return new InMemoryEmbeddingStore<>();
    }

    @Bean
    public AiCodeGeneratorService aiCodeGeneratorService(@Qualifier("DeepSeek") ChatModel chatModel, @Qualifier("DeepSeekStreaming") StreamingChatModel streamingChatModel){
        return AiServices.builder(AiCodeGeneratorService.class)
                .chatModel(chatModel)
                .streamingChatModel(streamingChatModel)
                .build();
    }


}
